import matplotlib.pyplot as plt
import torch.nn
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST

"""
获得MNIST数据集  28*28的手写数字图片
"""
def get_data_loader(is_train:bool):
    to_tensor  = transforms.Compose([transforms.ToTensor()])   # 数据类型，多维数组，张量
    data_set  = MNIST("",is_train,transform=to_tensor,download=True)# root 下载到根目录
    return DataLoader(data_set,batch_size=15,shuffle=True) # batch_size代表数据分批次训练，shuffle：随机打乱数据

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        #设置四个隐藏层 的 in 和 out
        self.fc1 = torch.nn.Linear(28*28,64)
        self.fc2 = torch.nn.Linear(64,64)
        self.fc3 = torch.nn.Linear(64,64)
        self.fc4 = torch.nn.Linear(64,10)

    def forward(self,x):
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))
        x = torch.nn.functional.relu(self.fc3(x))
        x = torch.nn.functional.log_softmax(self.fc4(x),1)
        return x

def evalute(test_data,net:Net):
    n_correct =0
    n_total = 0
    with torch.no_grad():
        for(x,y) in test_data:
            outputs = net.forward(x.view(-1,28*28))
            for i,output in enumerate(outputs):
                if torch.argmax(output) == y[i]:
                    n_correct+=1
                n_total+=1

    return n_correct/n_total

if __name__ == '__main__':
    train_data = get_data_loader(is_train=True)
    test_data = get_data_loader(is_train=False)
    net = Net()
    print("开始的正确率：",evalute(test_data,net))
    optimizer = torch.optim.Adam(net.parameters(),lr=0.001)
    for epoch in range(2):
        for (x,y) in train_data:
            net.zero_grad()
            output = net.forward(x.view(-1,28*28))
            loss = torch.nn.functional.nll_loss(output,y)
            loss.backward()
            optimizer.step()
        print("epoch",epoch,"准确率：",evalute((test_data,net)))

    for (n,(x,_)) in enumerate(test_data):
        if n>3:
            break
        predict = torch.argmax(net.forward(x[0].view(-1,28*28)))
        plt.figure(n)
        plt.imshow(x[0].view(28,28))
        plt.title("prediction:"+str(int(predict)))
    plt.show()

